import os
import sys
import subprocess
import shutil

from Data.DataStorage import HDFS
from Util.Process import INFO
from Util.Image import tiff_to_other_data, split_tile, filter_tile, merge_tile, resize_image
from Util.Return import get_config, get_path, get_dataset


class CWSAMAPI:
    def __init__(self):
        self.name = get_config("model", "CWSAM")['name']
        self.run_path = get_path(os.path.abspath(__file__), 4, get_config("model", "CWSAM")["run_path"])
        self.config_path = get_path(os.path.abspath(__file__), 4, get_config("model", "CWSAM")["config_path"])
        self.model_pth_path = get_path(os.path.abspath(__file__), 4, get_config("model", "CWSAM")["model_pth_path"])
        self.default_use_gpu = get_config("model", "CWSAM")['default_use_gpu']
        self.scene_id = None
        self.dataset = None
        self.tile_range = None
        self.file_list = []
        ###################################################################################
        self.temp_path = None
        self.raw_path = None
        self.input_path = None
        self.output_path = None
        self.result_path = None

    def initialize(self, data_list, info_list=None):  # 初始化
        self.scene_id = data_list[0]
        self.dataset = get_dataset(self.scene_id)
        self.tile_range = info_list[1]
        # PATH
        self.temp_path = os.path.join(get_path(os.path.abspath(__file__), 4, get_config("data", "TASK")['temp_path']),
                                      self.scene_id)
        self.raw_path = os.path.join(get_path(os.path.abspath(__file__), 4, get_config("data", "TASK")['raw_path']),
                                     self.scene_id + ".jpg")
        self.input_path = os.path.join(self.temp_path, "input_path")
        self.output_path = os.path.join(self.temp_path, "output_path")
        self.result_path = os.path.join(
            get_path(os.path.abspath(__file__), 4, get_config("data", "TASK")['result_path']), info_list[0] + ".jpg")

    def preprocess(self):  # 预处理: 拉取并处理成模型可输入的数据
        hdfs_available_path = "/" + get_config("data", self.dataset)['available_path']
        local_available_path = get_path(os.path.abspath(__file__), 4,
                                        get_config("data", self.dataset)['available_path'])
        hdfs = HDFS()
        flag, type = hdfs.exist(hdfs_available_path + "/" + self.scene_id + ".TIFF")
        if not flag:
            INFO("ERROR", f"CWSAM API preprocess: PULL AVAILABLE DATA FILE: Don't have {self.scene_id}.TIFF")
            raise Exception(f"CWSAM API preprocess: Don't have {self.scene_id}.TIFF")
        else:
            INFO("INFO", f"CWSAM API preprocess: PULL AVAILABLE DATA  {self.scene_id}.TIFF to local")
            hdfs.pull(hdfs_available_path + "/" + self.scene_id + ".TIFF", local_available_path)
            local_available_path = os.path.join(local_available_path, self.scene_id + ".TIFF")
            self.file_list.append(local_available_path)
            tiff_to_other_data(local_available_path, self.raw_path, False)
            self.file_list.append(self.raw_path)
            split_tile(self.raw_path, self.input_path)
            self.file_list.append(self.temp_path)
            filter_tile(self.input_path, self.tile_range)

    def request(self):  # 调用请求模型
        INFO("INFO",
             f"Request {self.name} Model {len([f for f in os.listdir(self.input_path) if os.path.isfile(os.path.join(self.input_path, f))])} images")
        args = ["--input_path", self.input_path,
                "--output_path", self.output_path,
                "--use_gpu_tag", self.default_use_gpu,
                "--config", self.config_path,
                "--model", self.model_pth_path]
        python_executable = sys.executable
        subprocess.run([python_executable, self.run_path] + args)

    def reprocess(self):  # 后处理: 将模型处理数据生成为最终结果
        merge_tile(self.output_path, self.result_path)  # 合并生成结果
        resize_image(self.result_path)
        hdfs = HDFS()
        hdfs.push("/" + get_config("data", "TASK")['result_path'], self.result_path)  # 上传结果
        for file_path in self.file_list:  # 清空产生的文件
            if os.path.isfile(file_path):
                os.remove(file_path)
            else:
                shutil.rmtree(file_path)
